"""
可视化数据转换模块
将栅格分析的可视化结果转换为GeoJSON格式的空间数据

主要功能：
1. 将掩膜、流向、流量累积等栅格数据转换为矢量格式
2. 提取河流网络、集水区边界等空间要素
3. 生成标准化的GeoJSON输出，包含完整元数据
4. 替代传统的图片输出，提供可交互的空间数据

输出文件：
- {outlet_id}_mask_boundary.geojson: 掩膜边界
- {outlet_id}_flow_directions.geojson: 流向矢量
- {outlet_id}_accumulation_points.geojson: 高流量累积点
- {outlet_id}_stream_network.geojson: 河流网络
- {outlet_id}_catchment_boundary.geojson: 集水区边界
- {outlet_id}_processing_points.geojson: 关键处理点位
"""

import os
import sys
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Any
import json

import numpy as np
import geopandas as gpd
import pandas as pd
from shapely.geometry import Point, LineString, Polygon, MultiPoint
from shapely.ops import unary_union
import rasterio
from rasterio.features import shapes
from loguru import logger

# 使用相对导入
from ..utils import config


class VisualizationGeoJSONConverter:
    """可视化数据GeoJSON转换器"""
    
    def __init__(self, output_dir: Optional[str] = None):
        """
        初始化转换器
        
        Args:
            output_dir: 输出目录，默认使用config中的设置
        """
        self.output_dir = output_dir or config.OUTPUT_DIR
        self.crs = "EPSG:4326"  # 标准坐标系
        
    def convert_mask_to_geojson(self, mask, catchment_poly, lat: float, lng: float, 
                               outlet_id: str) -> str:
        """
        将掩膜数据转换为GeoJSON格式
        
        Args:
            mask: 掩膜栅格数据
            catchment_poly: 集水区多边形
            lat: 纬度
            lng: 经度
            outlet_id: 出口点ID
            
        Returns:
            输出文件路径
        """
        
        logger.info(f"转换掩膜数据为GeoJSON: {outlet_id}")
        
        # 创建要素列表
        features = []
        
        # 1. 添加掩膜边界（从集水区多边形）
        if catchment_poly:
            features.append({
                'type': 'Feature',
                'geometry': catchment_poly.__geo_interface__,
                'properties': {
                    'feature_type': 'mask_boundary',
                    'outlet_id': outlet_id,
                    'description': 'Terminal unit catchment boundary',
                    'area_m2': catchment_poly.area if hasattr(catchment_poly, 'area') else None
                }
            })
        
        # 2. 添加出口点
        outlet_point = Point(lng, lat)
        features.append({
            'type': 'Feature',
            'geometry': outlet_point.__geo_interface__,
            'properties': {
                'feature_type': 'outlet_point',
                'outlet_id': outlet_id,
                'latitude': lat,
                'longitude': lng,
                'description': 'Watershed outlet point'
            }
        })
        
        # 3. 如果掩膜有栅格数据，提取边界
        if hasattr(mask, 'values') and hasattr(mask, 'extent'):
            try:
                # 将栅格转换为多边形
                mask_array = mask.values.astype(np.uint8)
                if mask_array.max() > 0:
                    # 使用rasterio提取形状
                    transform = rasterio.transform.from_bounds(
                        *mask.extent, mask_array.shape[1], mask_array.shape[0]
                    )
                    
                    mask_shapes = list(shapes(mask_array, transform=transform))
                    for geom, value in mask_shapes:
                        if value > 0:  # 只保留掩膜区域
                            features.append({
                                'type': 'Feature',
                                'geometry': geom,
                                'properties': {
                                    'feature_type': 'mask_raster',
                                    'outlet_id': outlet_id,
                                    'mask_value': int(value),
                                    'description': 'Raster mask polygon'
                                }
                            })
            except Exception as e:
                logger.warning(f"无法提取掩膜栅格形状: {str(e)}")
        
        # 创建GeoJSON
        geojson_data = {
            'type': 'FeatureCollection',
            'crs': {'type': 'name', 'properties': {'name': self.crs}},
            'features': features,
            'metadata': {
                'title': f'Watershed Mask Data - {outlet_id}',
                'description': 'Mask boundary and outlet point for watershed delineation',
                'outlet_id': outlet_id,
                'coordinate_system': self.crs,
                'created_by': 'mghydro visualization converter',
                'created_at': datetime.now().isoformat(),
                'feature_count': len(features),
                'data_type': 'mask_visualization'
            }
        }
        
        # 保存文件
        output_file = os.path.join(self.output_dir, f"{outlet_id}_mask_boundary.geojson")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(geojson_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"掩膜GeoJSON已保存: {output_file}")
        return output_file
    
    def convert_flow_directions_to_geojson(self, fdir, lat: float, lng: float, 
                                         outlet_id: str, dirmap: Dict, 
                                         catchment_poly=None) -> str:
        """
        将流向数据转换为GeoJSON格式
        
        Args:
            fdir: 流向栅格数据
            lat: 纬度
            lng: 经度
            outlet_id: 出口点ID
            dirmap: 流向映射字典
            catchment_poly: 集水区多边形（可选）
            
        Returns:
            输出文件路径
        """
        
        logger.info(f"转换流向数据为GeoJSON: {outlet_id}")
        
        features = []
        
        # 1. 添加出口点
        outlet_point = Point(lng, lat)
        features.append({
            'type': 'Feature',
            'geometry': outlet_point.__geo_interface__,
            'properties': {
                'feature_type': 'outlet_point',
                'outlet_id': outlet_id,
                'latitude': lat,
                'longitude': lng,
                'description': 'Watershed outlet point'
            }
        })
        
        # 2. 添加集水区边界（如果提供）
        if catchment_poly:
            features.append({
                'type': 'Feature',
                'geometry': catchment_poly.__geo_interface__,
                'properties': {
                    'feature_type': 'catchment_boundary',
                    'outlet_id': outlet_id,
                    'description': 'Catchment boundary for flow direction analysis'
                }
            })
        
        # 3. 提取流向统计信息
        if hasattr(fdir, 'values'):
            try:
                fdir_array = fdir.values
                unique_dirs, counts = np.unique(fdir_array[~np.isnan(fdir_array)], return_counts=True)
                
                # 创建流向统计点
                for direction, count in zip(unique_dirs, counts):
                    if direction in dirmap:
                        # 在集水区内找一个代表性位置
                        mask = fdir_array == direction
                        if np.any(mask):
                            y_indices, x_indices = np.where(mask)
                            if len(y_indices) > 0:
                                # 选择中心位置
                                center_idx = len(y_indices) // 2
                                y_idx, x_idx = y_indices[center_idx], x_indices[center_idx]
                                
                                # 转换为地理坐标
                                if hasattr(fdir, 'extent'):
                                    extent = fdir.extent
                                    x_coord = extent[0] + (x_idx / fdir_array.shape[1]) * (extent[1] - extent[0])
                                    y_coord = extent[2] + ((fdir_array.shape[0] - y_idx) / fdir_array.shape[0]) * (extent[3] - extent[2])
                                    
                                    point = Point(x_coord, y_coord)
                                    features.append({
                                        'type': 'Feature',
                                        'geometry': point.__geo_interface__,
                                        'properties': {
                                            'feature_type': 'flow_direction_sample',
                                            'outlet_id': outlet_id,
                                            'direction_code': int(direction),
                                            'direction_name': dirmap.get(direction, 'Unknown'),
                                            'pixel_count': int(count),
                                            'description': f'Sample point for flow direction {direction}'
                                        }
                                    })
            except Exception as e:
                logger.warning(f"无法提取流向统计信息: {str(e)}")
        
        # 创建GeoJSON
        geojson_data = {
            'type': 'FeatureCollection',
            'crs': {'type': 'name', 'properties': {'name': self.crs}},
            'features': features,
            'metadata': {
                'title': f'Flow Direction Data - {outlet_id}',
                'description': 'Flow direction analysis results with sample points',
                'outlet_id': outlet_id,
                'coordinate_system': self.crs,
                'created_by': 'mghydro visualization converter',
                'created_at': datetime.now().isoformat(),
                'feature_count': len(features),
                'data_type': 'flow_direction_visualization',
                'direction_mapping': dirmap
            }
        }
        
        # 保存文件
        output_file = os.path.join(self.output_dir, f"{outlet_id}_flow_directions.geojson")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(geojson_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"流向GeoJSON已保存: {output_file}")
        return output_file
    
    def convert_accumulation_to_geojson(self, acc, lat: float, lng: float, 
                                      lat_snap: float, lng_snap: float, 
                                      outlet_id: str, catchment_poly=None) -> str:
        """
        将流量累积数据转换为GeoJSON格式
        
        Args:
            acc: 流量累积栅格数据
            lat: 原始纬度
            lng: 原始经度
            lat_snap: 捕捉后纬度
            lng_snap: 捕捉后经度
            outlet_id: 出口点ID
            catchment_poly: 集水区多边形（可选）
            
        Returns:
            输出文件路径
        """
        
        logger.info(f"转换流量累积数据为GeoJSON: {outlet_id}")
        
        features = []
        
        # 1. 添加原始出口点
        original_point = Point(lng, lat)
        features.append({
            'type': 'Feature',
            'geometry': original_point.__geo_interface__,
            'properties': {
                'feature_type': 'original_outlet',
                'outlet_id': outlet_id,
                'latitude': lat,
                'longitude': lng,
                'description': 'Original outlet coordinates'
            }
        })
        
        # 2. 添加捕捉后出口点
        snapped_point = Point(lng_snap, lat_snap)
        features.append({
            'type': 'Feature',
            'geometry': snapped_point.__geo_interface__,
            'properties': {
                'feature_type': 'snapped_outlet',
                'outlet_id': outlet_id,
                'latitude': lat_snap,
                'longitude': lng_snap,
                'description': 'Snapped outlet coordinates on high accumulation area'
            }
        })
        
        # 3. 添加集水区边界
        if catchment_poly:
            features.append({
                'type': 'Feature',
                'geometry': catchment_poly.__geo_interface__,
                'properties': {
                    'feature_type': 'catchment_boundary',
                    'outlet_id': outlet_id,
                    'description': 'Catchment boundary for accumulation analysis'
                }
            })
        
        # 4. 提取高流量累积点
        if hasattr(acc, 'values') and hasattr(acc, 'extent'):
            try:
                acc_array = acc.values
                # 找到高累积值点（前5%）
                valid_acc = acc_array[~np.isnan(acc_array)]
                if len(valid_acc) > 0:
                    threshold = np.percentile(valid_acc, 95)  # 前5%的高值
                    high_acc_mask = acc_array >= threshold
                    
                    y_indices, x_indices = np.where(high_acc_mask)
                    extent = acc.extent
                    
                    # 限制点数量，避免过多
                    max_points = 50
                    if len(y_indices) > max_points:
                        step = len(y_indices) // max_points
                        y_indices = y_indices[::step]
                        x_indices = x_indices[::step]
                    
                    for y_idx, x_idx in zip(y_indices, x_indices):
                        x_coord = extent[0] + (x_idx / acc_array.shape[1]) * (extent[1] - extent[0])
                        y_coord = extent[2] + ((acc_array.shape[0] - y_idx) / acc_array.shape[0]) * (extent[3] - extent[2])
                        
                        acc_value = acc_array[y_idx, x_idx]
                        
                        point = Point(x_coord, y_coord)
                        features.append({
                            'type': 'Feature',
                            'geometry': point.__geo_interface__,
                            'properties': {
                                'feature_type': 'high_accumulation_point',
                                'outlet_id': outlet_id,
                                'accumulation_value': float(acc_value),
                                'percentile_rank': 95,
                                'description': 'High flow accumulation point (top 5%)'
                            }
                        })
            except Exception as e:
                logger.warning(f"无法提取高流量累积点: {str(e)}")
        
        # 创建GeoJSON
        geojson_data = {
            'type': 'FeatureCollection',
            'crs': {'type': 'name', 'properties': {'name': self.crs}},
            'features': features,
            'metadata': {
                'title': f'Flow Accumulation Data - {outlet_id}',
                'description': 'Flow accumulation analysis with outlet snapping results',
                'outlet_id': outlet_id,
                'coordinate_system': self.crs,
                'created_by': 'mghydro visualization converter',
                'created_at': datetime.now().isoformat(),
                'feature_count': len(features),
                'data_type': 'accumulation_visualization',
                'snapping_info': {
                    'original_coords': [lng, lat],
                    'snapped_coords': [lng_snap, lat_snap],
                    'displacement_m': Point(lng, lat).distance(Point(lng_snap, lat_snap)) * 111000  # 近似转换为米
                }
            }
        }
        
        # 保存文件
        output_file = os.path.join(self.output_dir, f"{outlet_id}_accumulation_points.geojson")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(geojson_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"流量累积GeoJSON已保存: {output_file}")
        return output_file
    
    def convert_streams_to_geojson(self, streams, catchment_poly, lat: float, lng: float,
                                 lat_snap: float, lng_snap: float, outlet_id: str, 
                                 numpixels: int) -> str:
        """
        将河流网络数据转换为GeoJSON格式
        
        Args:
            streams: 河流栅格数据
            catchment_poly: 集水区多边形
            lat: 原始纬度
            lng: 原始经度
            lat_snap: 捕捉后纬度
            lng_snap: 捕捉后经度
            outlet_id: 出口点ID
            numpixels: 像素阈值
            
        Returns:
            输出文件路径
        """
        
        logger.info(f"转换河流网络数据为GeoJSON: {outlet_id}")
        
        features = []
        
        # 1. 添加出口点
        features.extend([
            {
                'type': 'Feature',
                'geometry': Point(lng, lat).__geo_interface__,
                'properties': {
                    'feature_type': 'original_outlet',
                    'outlet_id': outlet_id,
                    'latitude': lat,
                    'longitude': lng,
                    'description': 'Original outlet coordinates'
                }
            },
            {
                'type': 'Feature',
                'geometry': Point(lng_snap, lat_snap).__geo_interface__,
                'properties': {
                    'feature_type': 'snapped_outlet',
                    'outlet_id': outlet_id,
                    'latitude': lat_snap,
                    'longitude': lng_snap,
                    'description': 'Snapped outlet coordinates'
                }
            }
        ])
        
        # 2. 添加集水区边界
        if catchment_poly:
            features.append({
                'type': 'Feature',
                'geometry': catchment_poly.__geo_interface__,
                'properties': {
                    'feature_type': 'catchment_boundary',
                    'outlet_id': outlet_id,
                    'description': 'Catchment boundary for stream network'
                }
            })
        
        # 3. 提取河流网络
        if hasattr(streams, 'values') and hasattr(streams, 'extent'):
            try:
                streams_array = streams.values
                # 提取河流像素位置
                stream_mask = streams_array > numpixels
                
                if np.any(stream_mask):
                    y_indices, x_indices = np.where(stream_mask)
                    extent = streams.extent
                    
                    # 将河流像素转换为点
                    stream_points = []
                    for y_idx, x_idx in zip(y_indices, x_indices):
                        x_coord = extent[0] + (x_idx / streams_array.shape[1]) * (extent[1] - extent[0])
                        y_coord = extent[2] + ((streams_array.shape[0] - y_idx) / streams_array.shape[0]) * (extent[3] - extent[2])
                        stream_points.append(Point(x_coord, y_coord))
                    
                    # 如果点太多，进行采样
                    if len(stream_points) > 1000:
                        step = len(stream_points) // 1000
                        stream_points = stream_points[::step]
                    
                    # 添加河流点要素
                    for i, point in enumerate(stream_points):
                        features.append({
                            'type': 'Feature',
                            'geometry': point.__geo_interface__,
                            'properties': {
                                'feature_type': 'stream_point',
                                'outlet_id': outlet_id,
                                'point_id': i,
                                'threshold_pixels': numpixels,
                                'description': f'Stream network point (threshold: {numpixels} pixels)'
                            }
                        })
                    
                    # 尝试创建河流线要素（简化版）
                    if len(stream_points) > 1:
                        # 创建一个简化的河流网络线
                        coords = [(p.x, p.y) for p in stream_points]
                        if len(coords) >= 2:
                            # 简单连接所有点（实际应用中需要更复杂的拓扑分析）
                            stream_line = LineString(coords)
                            features.append({
                                'type': 'Feature',
                                'geometry': stream_line.__geo_interface__,
                                'properties': {
                                    'feature_type': 'stream_network',
                                    'outlet_id': outlet_id,
                                    'point_count': len(coords),
                                    'threshold_pixels': numpixels,
                                    'description': 'Simplified stream network line'
                                }
                            })
            except Exception as e:
                logger.warning(f"无法提取河流网络: {str(e)}")
        
        # 创建GeoJSON
        geojson_data = {
            'type': 'FeatureCollection',
            'crs': {'type': 'name', 'properties': {'name': self.crs}},
            'features': features,
            'metadata': {
                'title': f'Stream Network Data - {outlet_id}',
                'description': 'Stream network extracted from flow accumulation',
                'outlet_id': outlet_id,
                'coordinate_system': self.crs,
                'created_by': 'mghydro visualization converter',
                'created_at': datetime.now().isoformat(),
                'feature_count': len(features),
                'data_type': 'stream_network_visualization',
                'extraction_parameters': {
                    'pixel_threshold': numpixels,
                    'description': f'Streams with upstream area > {numpixels} pixels'
                }
            }
        }
        
        # 保存文件
        output_file = os.path.join(self.output_dir, f"{outlet_id}_stream_network.geojson")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(geojson_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"河流网络GeoJSON已保存: {output_file}")
        return output_file
    
    def create_processing_summary(self, outlet_id: str, generated_files: List[str]) -> str:
        """
        创建处理摘要GeoJSON文件
        
        Args:
            outlet_id: 出口点ID
            generated_files: 生成的文件列表
            
        Returns:
            摘要文件路径
        """
        
        # 创建摘要数据
        summary_data = {
            'type': 'FeatureCollection',
            'crs': {'type': 'name', 'properties': {'name': self.crs}},
            'features': [],  # 摘要不包含空间要素
            'metadata': {
                'title': f'Visualization Processing Summary - {outlet_id}',
                'description': 'Summary of visualization data conversion to GeoJSON format',
                'outlet_id': outlet_id,
                'coordinate_system': self.crs,
                'created_by': 'mghydro visualization converter',
                'created_at': datetime.now().isoformat(),
                'processing_type': 'visualization_to_geojson',
                'generated_files': generated_files,
                'file_count': len(generated_files),
                'conversion_info': {
                    'input_format': 'raster_visualization_data',
                    'output_format': 'geojson',
                    'spatial_reference': self.crs,
                    'feature_types': [
                        'mask_boundary',
                        'flow_directions', 
                        'accumulation_points',
                        'stream_network'
                    ]
                }
            }
        }
        
        # 保存摘要文件
        summary_file = os.path.join(self.output_dir, f"{outlet_id}_visualization_summary.geojson")
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(summary_data, f, indent=2, ensure_ascii=False)
        
        logger.info(f"可视化处理摘要已保存: {summary_file}")
        return summary_file


def main():
    """主函数，演示可视化数据转换功能"""
    
    logger.info("=== 可视化数据GeoJSON转换工具演示 ===")
    
    # 创建转换器
    converter = VisualizationGeoJSONConverter()
    
    # 模拟数据（实际使用中会从栅格分析获得）
    outlet_id = "demo_6401070"
    lat, lng = 64.71072, -21.60337
    lat_snap, lng_snap = 64.71080, -21.60340
    
    logger.info("可视化数据GeoJSON转换工具已准备就绪")
    logger.info("使用示例:")
    logger.info("  converter = VisualizationGeoJSONConverter()")
    logger.info("  converter.convert_mask_to_geojson(mask, poly, lat, lng, outlet_id)")
    logger.info("  converter.convert_accumulation_to_geojson(acc, lat, lng, lat_snap, lng_snap, outlet_id)")


if __name__ == "__main__":
    main()